"""
默认配置参数
"""

import os
from typing import Dict, Any, Optional


class BaseConfig:
    """基础配置类"""
    
    def __init__(self):
        # 数据配置
        self.data_dir = "./data"
        self.batch_size = 4
        self.num_workers = 4
        self.pin_memory = True
        self.persistent_workers = True
        
        # 模型配置
        self.model_type = "mhd"  # mhd 或 maxwell
        self.input_dim = 6  # 输入场分量数
        self.output_dim = 6  # 输出场分量数
        self.modes = 12  # 傅里叶模式数
        self.width = 32  # 隐藏层宽度
        self.n_layers = 4  # 网络层数
        
        # 训练配置
        self.epochs = 100
        self.lr = 1e-3
        self.weight_decay = 1e-5
        self.scheduler_patience = 10
        self.scheduler_factor = 0.5
        self.early_stopping_patience = 20
        
        # 损失函数配置
        self.loss_weights = {
            "data": 1.0,
            "ic": 1.0,
            "pde": 0.1,
            "constraint": 0.1
        }
        
        # 设备配置
        self.device = "auto"  # auto, cpu, cuda
        
        # 输出配置
        self.output_dir = "./outputs"
        self.checkpoint_dir = "./checkpoints"
        self.log_dir = "./logs"
        self.plot_dir = "./plots"
        
        # 其他配置
        self.seed = 42
        self.log_interval = 10
        self.eval_interval = 5
        self.save_interval = 20
        
    def to_dict(self) -> Dict[str, Any]:
        """转换为字典"""
        return {k: v for k, v in self.__dict__.items()}


class MHDConfig(BaseConfig):
    """MHD模型配置"""
    
    def __init__(self):
        super().__init__()
        
        # MHD特定配置
        self.model_type = "mhd"
        self.input_dim = 6  # [rho, vx, vy, vz, Bx, By, Bz] - 7个分量，但通常去掉密度或只使用6个
        self.output_dim = 6  # 预测6个分量
        
        # MHD物理参数
        self.viscosity = 0.01
        self.resistivity = 0.01
        self.gamma = 5/3  # 绝热指数
        
        # MHD数据配置
        self.time_steps = 10
        self.spatial_dim = [32, 32, 32]  # 3D网格尺寸
        self.normalize_fields = True
        
        # MHD损失权重
        self.loss_weights = {
            "data": 1.0,
            "ic": 1.0,
            "pde": 0.1,
            "constraint": 0.1,
            "continuity": 0.1,
            "momentum": 0.1,
            "induction": 0.1
        }


class MaxwellConfig(BaseConfig):
    """Maxwell方程模型配置"""
    
    def __init__(self):
        super().__init__()
        
        # Maxwell特定配置
        self.model_type = "maxwell"
        self.input_dim = 6  # [Ex, Ey, Ez, Bx, By, Bz]
        self.output_dim = 6  # 预测6个分量
        
        # Maxwell物理参数
        self.c = 1.0  # 光速
        self.epsilon_0 = 1.0  # 真空介电常数
        self.mu_0 = 1.0  # 真空磁导率
        
        # Maxwell数据配置
        self.time_steps = 10
        self.spatial_dim = [32, 32, 32]  # 3D网格尺寸
        self.normalize_fields = True
        
        # Maxwell损失权重
        self.loss_weights = {
            "data": 1.0,
            "ic": 1.0,
            "pde": 0.1,
            "constraint": 0.1,
            "faraday": 0.1,
            "ampere": 0.1,
            "gauss_e": 0.1,
            "gauss_b": 0.1
        }


def get_config(model_type: str, **kwargs) -> BaseConfig:
    """获取配置对象
    
    参数:
        model_type: 模型类型 ('mhd' 或 'maxwell')
        **kwargs: 其他配置参数
        
    返回:
        配置对象
    """
    if model_type.lower() == "mhd":
        config = MHDConfig()
    elif model_type.lower() == "maxwell":
        config = MaxwellConfig()
    else:
        raise ValueError(f"不支持的模型类型: {model_type}")
    
    # 更新配置
    for key, value in kwargs.items():
        if hasattr(config, key):
            setattr(config, key, value)
        else:
            print(f"警告: 配置中没有参数 {key}")
    
    return config


def save_config(config: BaseConfig, save_path: str):
    """保存配置到文件
    
    参数:
        config: 配置对象
        save_path: 保存路径
    """
    os.makedirs(os.path.dirname(save_path), exist_ok=True)
    
    with open(save_path, 'w') as f:
        for key, value in config.to_dict().items():
            f.write(f"{key}: {value}\n")


def load_config(load_path: str, model_type: str) -> BaseConfig:
    """从文件加载配置
    
    参数:
        load_path: 加载路径
        model_type: 模型类型
        
    返回:
        配置对象
    """
    config = get_config(model_type)
    
    with open(load_path, 'r') as f:
        for line in f:
            if ':' in line:
                key, value = line.strip().split(':', 1)
                if hasattr(config, key):
                    # 尝试转换为适当的类型
                    try:
                        if value.lower() == 'true':
                            setattr(config, key, True)
                        elif value.lower() == 'false':
                            setattr(config, key, False)
                        elif '.' in value:
                            setattr(config, key, float(value))
                        else:
                            setattr(config, key, int(value))
                    except ValueError:
                        setattr(config, key, value)
    
    return config